from sqlalchemy import Enum
from pydantic import BaseModel, ValidationError, constr, Field
import json
import logging


from queue import Queue


TASK_QUEUE = Queue(20)

LOGGER = logging.getLogger(__name__)


# 走redis 发布的消息有两种，1生成，2退出
class TaskType:
    GEN_THESIS: int = 1
    STOP_LISTEN: int = 11


class TaskModel(BaseModel):
    type: int = TaskType.GEN_THESIS
    thesisId: int = -1


def popTaskBlocking():
    try:
        return TASK_QUEUE.get(timeout=5)
    except Exception as e:
        LOGGER.info("任务队列为空，没有文章生成任务")
        return None


def pubTask(task: TaskModel):
    TASK_QUEUE.put(json.dumps(task.dict()))
    LOGGER.info(f"发布一个任务 {task}")


def pubGenerateThesisTask(thesisId: int):
    task = TaskModel(type=TaskType.GEN_THESIS, thesisId=thesisId)
    pubTask(task)
